from fastapi import FastAPI, File, UploadFile
import shutil
import os
import cv2
import numpy as np
import pandas as pd
import logging
import uvicorn
from skimage.measure import label, regionprops
from fastapi.middleware.cors import CORSMiddleware
from time import time
from typing import List
from fastapi.staticfiles import StaticFiles
import requests

import concurrent.futures
app = FastAPI()
app.mount("/static", StaticFiles(directory="OCT_layer_cal/image/processed_results"), name="static")
app.mount("/colorinput", StaticFiles(directory="user_color_data"), name="colorinput")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源跨域
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ImagePath:
    def __init__(self, path: str):
        self.path = path

logging.basicConfig(level=logging.DEBUG)


def extract_center_line(binary_image):
    # 使用传统的细化算法进行中心线提取
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    done = False
    size = np.size(binary_image)
    skel = np.zeros(binary_image.shape, np.uint8)

    while not done:
        eroded = cv2.erode(binary_image, element)
        temp = cv2.dilate(eroded, element)
        temp = cv2.subtract(binary_image, temp)
        skel = cv2.bitwise_or(skel, temp)
        binary_image = eroded.copy()

        zeros = size - cv2.countNonZero(binary_image)
        if zeros == size:
            done = True

    return skel

def calculate_average_width(center_line_points, contours):
    average_widths = []

    for point in center_line_points:
        distances = []
        for contour in contours:
            distance = cv2.pointPolygonTest(contour, (int(point[1]), int(point[0])), True)
            distances.append(abs(distance))

        # 计算中心线到轮廓的最小距离
        min_distance = min(distances)

        # 将最小距离乘以2，即为连通域在该点的宽度
        width_at_point = min_distance * 2
        average_widths.append(width_at_point)

    # 计算平均宽度
    average_width = np.mean(average_widths)
    return average_width


def process_image(image_path, output_path_csv, color, slice):
    # 读取图像并进行二值化处理
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    _, binary_image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

    # 提取中心线
    center_line = extract_center_line(binary_image)
    # _, contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 本地必须用下面这个
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)


    # 计算平均宽度
    center_line_points = np.column_stack(np.where(center_line > 0))
    average_width = calculate_average_width(center_line_points, contours)

    # 初始化layer变量为默认值
    layer = 'Unknown Layer'

    # 根据颜色值设置layer
    if all(color == [0, 0, 255]):
        layer = 'ILM-IPL'
    elif all(color == [0, 255, 0]):
        layer = 'IPL-OPL'
    elif all(color == [0, 255, 192]):
        layer = 'OPL-ISOS'
    elif all(color == [156, 255, 0]):
        layer = 'ISOS-RPE'
    elif all(color == [255, 0, 0]):
        layer = 'RPE-BM'

    data = {'Slice': [slice], 'Layer': [layer], 'Average Height (px)': [average_width]}
    df = pd.DataFrame(data)

    # 获取输出文件夹名和csv文件名
    folder_name = output_path_csv.split('/')[-1]
    csv_file_name = f"{folder_name}.csv"

    # 构建完整的输出CSV路径
    full_output_path = f"{output_path_csv}/{csv_file_name}"
    # 保证文件夹一定存在
    os.makedirs(os.path.dirname(full_output_path), exist_ok=True)
    df.to_csv(full_output_path, mode='a', index=False, header=(not os.path.isfile(full_output_path)))

def process_color_label_take(image_path, output_path, folder_name, slice_image_name):
    # 处理当前图像
    color_label_take(image_path, output_path, folder_name, slice_image_name)
    print(f"Image {slice_image_name} processed successfully")

@app.post("/upload_image/")
async def upload_image(image: UploadFile = File(...)):
    timestamp = int(time())
    os.makedirs(f"{os.getcwd()}/user_color_data/{timestamp}", exist_ok=True)
    save_path = f"{os.getcwd()}/user_color_data/{timestamp}/{image.filename}"
    with open(save_path, "wb") as img:
        shutil.copyfileobj(image.file, img)

    return {"message": "Image uploaded successfully","time":timestamp}



@app.post("/process_image/")
async def process_uploaded_image(user_string: int):
    input_path = f'{os.getcwd()}/user_color_data/{user_string}'
    # print(input_path)

    # cal
    upload_url = 'http://0.0.0.0:30002/run/optovue/five/opti/'

    # image_filename = '/home/imed/Desktop/oct5-thick-layer-analysis/server/user_color_data/1713251553/image_0.png'
    output_path1 = f'{os.getcwd()}/user_color_data/{user_string}/{user_string}/raw'
    output_path2 = f'{os.getcwd()}/user_color_data/{user_string}/{user_string}/color'
    os.makedirs(output_path1, exist_ok=True)
    os.makedirs(output_path2, exist_ok=True)

    # 遍历原始图像文件夹
    for image_file in os.listdir(input_path):
        if image_file.endswith('.png') or image_file.endswith('.jpg'):
            # 构建完整的输入图像路径
            image_path = os.path.join(input_path, image_file)
            files = {'file': open(image_path, 'rb')}
            response = requests.post(upload_url, files=files)
            with open(f'{output_path1}/raw.png', 'wb') as r_f:
                r_f.write(response.content)

            color_list = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (192, 255, 0), (0, 255, 156)]

            ori_img = cv2.imread(f'{output_path1}/raw.png')
            zero_img = np.zeros_like(ori_img, dtype=np.uint8)
            ori_img = ori_img[:, :, 0]

            for i in range(1, 6):
                zero_img[np.where(ori_img == i)] = color_list[i - 1]

            cv2.imwrite(f'{output_path2}/color.png', zero_img)


    # statis


    # input_path1 = f'{os.getcwd()}/user_color_data/{user_string}/{user_string}'
    # print(input_path)
    output_path = f'{os.getcwd()}/OCT_layer_cal'
    # print(output_path)
    os.makedirs(output_path, exist_ok=True)

    # 遍历原始图像文件夹
    image_files = [f for f in os.listdir(output_path2) if f.endswith('.png') or f.endswith('.jpg')]
    folder_name = 'processed_results'
    slice_image_name = f'{user_string}'

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = []
        for image_file in image_files:
            # 构建完整的输入图像路径
            image_path = os.path.join(output_path2, image_file)

            # 提交任务到线程池
            future = executor.submit(process_color_label_take, image_path, output_path, folder_name, slice_image_name)
            futures.append(future)

        # 等待所有任务完成
        concurrent.futures.wait(futures)

    return {"processed_image": f"Image {os.listdir(output_path2)} processed successfully"}


@app.get("/list_images/")
async def list_images(user_string: int):
    image_paths = []
    folder_path = f'{os.getcwd()}/OCT_layer_cal/image/processed_results/{user_string}/'
    # 遍历指定文件夹中的所有文件
    for file_name in os.listdir(folder_path):
        file_path = os.path.join(f'http://114.55.245.149:7833/static/{user_string}/', file_name)

        # 检查文件是否为图片文件（这里仅检查扩展名）
        if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
            image_paths.append(file_path)

    return image_paths

@app.post("/mix_together/")
async def mix_together(user_string: int,base_image_name: str):
    # 读取两张图片
    imagepath1 = f"{os.getcwd()}/user_color_data/{user_string}/{base_image_name}"
    imagepath2 = f"{os.getcwd()}/user_color_data/{user_string}/{user_string}/color/color.png"
    print(imagepath1)
    print(imagepath2)

    image1 = cv2.imread(imagepath1)
    image2 = cv2.imread(imagepath2)

    if image1 is not None and image2 is not None:
        # 将第一张图片缩放到和第二张图片相同大小
        # image1_resized = cv2.resize(image1, (image2.shape[1], image2.shape[0]))

        # 将两张图片融合
        blended_image = cv2.addWeighted(image1, 0.5, image2, 0.5, 0)

        # 保存融合后的图片
        cv2.imwrite(f"{os.getcwd()}/user_color_data/{user_string}/blended_image.png", blended_image)
    else:
        print("无法读取图片，请检查文件路径是否正确。")

    return {"message":"complete"}

@app.post("/mix_images/")
async def mix_images(user_string: int,base_image_name: str):
    mix_image_name = []
    # 加载主图片
    main_image_path = f"{os.getcwd()}/user_color_data/{user_string}/{base_image_name}"
    # 目标文件夹里的所有图片
    image_folder_path = f'{os.getcwd()}/OCT_layer_cal/image/processed_results/{user_string}'
    # 创建新文件夹保存融合后的图片
    os.makedirs(f'{os.getcwd()}/user_color_data/{user_string}/output_mixfolder', exist_ok=True)

    main_image = cv2.imread(main_image_path)

    # 循环处理每张图片
    for image_file in os.listdir(image_folder_path):
        # 加载目标文件夹里的图片
        image_path = os.path.join(image_folder_path, image_file)
        print(image_path)
        target_image = cv2.imread(image_path)
        print(target_image)
        if target_image is None or target_image.size == 0:
            continue  # 跳过空图片

        if target_image.shape[:2] == (0, 0):
            continue  # 跳过尺寸为空的图片

        # 调整目标图片大小与主图片相匹配
        target_image = cv2.resize(target_image, (main_image.shape[1], main_image.shape[0]))

        # 在主图片上叠加目标图片（这里以简单的叠加为例）
        blended_image = cv2.addWeighted(main_image, 0.5, target_image, 0.5, 0)

        new_file_name = "mix" + image_file
        mix_image_name.append(new_file_name)
        # 保存融合后的图片到新文件夹中
        output_path = os.path.join(f'{os.getcwd()}/user_color_data/{user_string}/output_mixfolder', new_file_name)

        cv2.imwrite(output_path, blended_image)

    return mix_image_name

def color_label_take(label, output_path, folder_name, slice_image_name):
    label = cv2.imread(label)
    # 获取图像形状
    height, width, _ = label.shape

    # 提取颜色信息
    unique_colors = np.unique(label.reshape(-1, label.shape[2]), axis=0)
    unique_colors = unique_colors[unique_colors.any(axis=1)]
    # 确保输出路径存在
    output_path_csv = os.path.join(output_path, 'csv', folder_name)
    output_path_image = os.path.join(output_path, 'image', folder_name, slice_image_name)
    os.makedirs(output_path_image, exist_ok=True)
    os.makedirs(output_path_csv, exist_ok=True)

    for color in unique_colors:
        # 创建一个与原始图像相同形状的零数组
        result_array_label = np.zeros((height, width, label.shape[2]), dtype=np.uint8)

        # 找出当前颜色的像素位置
        color_mask = (label == color).all(axis=-1)

        # 对于3通道图像，进行连通域分析（需要先转为灰度图）
        gray_color_mask = color_mask.sum(axis=-1) > 0
        _, cc_labels, stats, _ = cv2.connectedComponentsWithStats(gray_color_mask.astype(np.uint8), connectivity=4)

        # 找到面积最大的连通域
        max_area = -1
        max_label = None
        for i, prop in enumerate(regionprops(cc_labels)):
            if prop.area > max_area and prop.area > 0:
                max_area = prop.area
                max_label = i + 1  # 连通域标签从1开始

        # 将最大连通域内的像素恢复成当前颜色
        if max_label is not None:
            result_array_label[color_mask & (cc_labels == max_label)] = color

        # 保存每个颜色的最大连通域结果
        cv2.imwrite(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg', result_array_label)

        process_image(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg', output_path_csv, color,
                      slice_image_name)
        # print(f'oct厚层分析完成---{slice_image_name}')


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)
